
#include "fun.h"
#include "hs_cuda_inputwise.cuh"

struct HistStatisFunctor
{

    HistStatisFunctor(int bin, float* p_hist)
        : m_bin(bin), m_p_hist(p_hist)
    { }

    __device__ void operator()(float val)
    {
        int rval = (int)rintf(val); // 四舍五入得到索引
        if (rval < m_bin) {
            // 原子加
            atomicAdd(&m_p_hist[rval], 1);
        }
    }

private:
    int m_bin;
    float* m_p_hist; // 输出直返图, 长度为m_bin
};

void histStatis(const hs::raster::RasterData<float> & mat, int bin, hs::raster::RasterData<float> & hist)
{
    /* 将数据拷贝到的GPU */
    hscuda::raster::RasterData<float> device_mat;
    hscuda::raster::host2device(mat.view(), device_mat);

    /* 输出分配内存 */
    hscuda::raster::RasterData<float> device_hist(bin, 1);

    /* 调用 inputwise */
    int N = device_mat.view().numel();
    hs::cuda::inputwise::Unary(N
        , device_mat.view().pData()
        , HistStatisFunctor(bin, device_hist.view().pData()));
    cudaDeviceSynchronize();
    
    /* 输出复制到主机端 */
    device_hist.toHost(hist);
}

struct PaddingZeroFunctor
{

    PaddingZeroFunctor(int dst_size_x, int dst_size_y, float* p_dst_data, int src_size_x, int src_size_y)
        : m_dst_size_x(dst_size_x), m_dst_size_y(dst_size_y), m_p_dst_data(p_dst_data), m_src_size_x(src_size_x), m_src_size_y(src_size_y)
    { }

    __device__ void operator()(int id, float val)
    {
        if (id < m_src_size_x * m_src_size_y) {
            // id 是输入图像的线性索引, 转换为平面索引
            int yid = id / m_src_size_x;
            int xid = id - yid * m_src_size_x;

            // 加上偏移量
            yid = yid + 1;
            xid = xid + 1;

            // 转换为输出图线性索引
            int dst_id = xid + yid * m_dst_size_x;
            m_p_dst_data[dst_id] = val;
        }
    }

private:
    int m_src_size_y;
    int m_src_size_x;
    int m_dst_size_x;
    int m_dst_size_y;
    float* m_p_dst_data;
};


void paddingZero(const hs::raster::RasterData<float> & src_mat, hs::raster::RasterData<float> & dst_mat)
{
    /* 将数据拷贝到GPU */
    hscuda::raster::RasterData<float> device_src_mat;
    hscuda::raster::host2device(src_mat.view(), device_src_mat);

    /* 输出数据分配内存 */
    int size_x = src_mat.view().sizeX();
    int size_y = src_mat.view().sizeY();
    hscuda::raster::RasterData<float> device_dst_mat(size_x + 2, size_y + 2);
    // 初始化为0
    cudaMemset(device_dst_mat.view().pData(), 0, sizeof(float) * device_dst_mat.view().numel());

    /* 调用 inputwiseIndex */
    int N = device_src_mat.view().numel();
    hs::cuda::inputwiseIndex::Unary(N
        , device_src_mat.view().pData()
        , PaddingZeroFunctor(
            device_dst_mat.view().sizeX()
            , device_dst_mat.view().sizeY()
            , device_dst_mat.view().pData()
            , size_x
            , size_y)
        );
    cudaDeviceSynchronize();

    /* 输出复制到主机端 */
    device_dst_mat.toHost(dst_mat);
}
